import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim
import os
import matplotlib.pyplot as plt
import tkinter as tk
from tkinter import filedialog
from datetime import datetime


class ImageComparator:
    def __init__(self, min_contour_area=50, morph_kernel_size=3, ssim_threshold=0.85,
                 max_features=5000, good_match_percent=0.2):
        """
        初始化图像比较器

        参数:
        min_contour_area: 最小轮廓面积阈值，用于过滤小差异
        morph_kernel_size: 形态学操作的核大小
        ssim_threshold: 结构相似性阈值，用于确定差异区域
        max_features: 特征检测的最大特征点数
        good_match_percent: 用于配准的良好匹配点的百分比
        """
        self.min_contour_area = min_contour_area
        self.morph_kernel_size = morph_kernel_size
        self.ssim_threshold = ssim_threshold
        self.max_features = max_features
        self.good_match_percent = good_match_percent
        self.last_result = None

    def preprocess_image(self, image):
        """图像预处理：灰度转换、直方图均衡化"""
        # 转换为灰度图
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()

        # 自适应直方图均衡化增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        equalized = clahe.apply(gray)

        # 高斯模糊减少噪声
        blurred = cv2.GaussianBlur(equalized, (5, 5), 0)
        return blurred

    def align_images(self, img1, img2):
        """使用特征匹配对齐两张图片"""
        sift = cv2.SIFT_create(self.max_features)
        kp1, des1 = sift.detectAndCompute(img1, None)
        kp2, des2 = sift.detectAndCompute(img2, None)

        if des1 is None or des2 is None or len(kp1) < 4 or len(kp2) < 4:
            return img1, img2, None

        # 使用FLANN匹配器进行特征匹配
        flann = cv2.FlannBasedMatcher(dict(algorithm=1, trees=5), dict(checks=50))
        matches = flann.knnMatch(des1, des2, k=2)

        good_matches = []
        for m, n in matches:
            if m.distance < 0.7 * n.distance:
                good_matches.append(m)

        if len(good_matches) < 4:
            return img1, img2, None

        # 获取匹配点的坐标
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)

        # 计算单应性矩阵
        M, _ = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 5.0)

        if M is not None:
            aligned = cv2.warpPerspective(img2, M, (img1.shape[1], img1.shape[0]))
            return img1, aligned, M

        return img1, img2, None

    def find_differences(self, img1, img2):
        """找出两张图片之间的差异"""
        if img1.shape != img2.shape:
            img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))

        # 预处理
        gray1 = self.preprocess_image(img1)
        gray2 = self.preprocess_image(img2)

        # 对齐图像
        aligned1, aligned2, M = self.align_images(gray1, gray2)

        # 使用结构相似性（SSIM）计算图像差异
        score, diff = ssim(aligned1, aligned2, full=True)
        diff = (diff * 255).astype("uint8")

        _, thresh = cv2.threshold(diff, int(255 * (1 - self.ssim_threshold)),
                                  255, cv2.THRESH_BINARY_INV)

        # 形态学操作增强差异区域
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
                                           (self.morph_kernel_size, self.morph_kernel_size))
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=3)
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=1)

        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        valid_contours = [contour for contour in contours if cv2.contourArea(contour) > self.min_contour_area]

        return valid_contours, thresh, M

    def visualize_differences(self, img1, img2, contours, M=None, output_path=None, show_result=False):
        """可视化差异结果：并排显示两张图片并在上面标记差异"""
        img1_display = img1.copy()
        img2_display = img2.copy()

        # 在第一张图上绘制差异区域
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            cv2.rectangle(img1_display, (x, y), (x + w, y + h), (0, 0, 255), 2)

        if M is not None:
            M_inv = np.linalg.inv(M)
            for contour in contours:
                contour_homogeneous = np.ones((contour.shape[0], contour.shape[1], 3))
                contour_homogeneous[:, :, :2] = contour
                transformed_contour = np.dot(contour_homogeneous, M_inv.T)
                transformed_contour = transformed_contour[:, :, :2].astype(np.int32)
                x, y, w, h = cv2.boundingRect(transformed_contour)
                cv2.rectangle(img2_display, (x, y), (x + w, y + h), (0, 0, 255), 2)
        else:
            for contour in contours:
                x, y, w, h = cv2.boundingRect(contour)
                cv2.rectangle(img2_display, (x, y), (x + w, y + h), (0, 0, 255), 2)

        height = max(img1_display.shape[0], img2_display.shape[0])
        width1 = int(img1_display.shape[1] * height / img1_display.shape[0])
        width2 = int(img2_display.shape[1] * height / img2_display.shape[0])

        img1_display = cv2.resize(img1_display, (width1, height))
        img2_display = cv2.resize(img2_display, (width2, height))

        combined = np.hstack((img1_display, img2_display))

        cv2.putText(combined, "Image 1", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        cv2.putText(combined, "Image 2", (width1 + 10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

        diff_text = f"Differences found: {len(contours)}"
        cv2.putText(combined, diff_text, (10, combined.shape[0] - 20),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)

        # 保存结果
        if output_path:
            cv2.imwrite(output_path, combined)

        # 显示结果
        if show_result:
            plt.figure(figsize=(20, 10))
            plt.imshow(cv2.cvtColor(combined, cv2.COLOR_BGR2RGB))
            plt.title("Image Comparison Results")
            plt.axis('off')
            plt.show()

        return combined

    def compare_images(self, img_path1, img_path2, output_path=None, show_result=False):
        """
        比较两张图片并找出差异

        参数:
        img_path1: 第一张图片路径
        img_path2: 第二张图片路径
        output_path: 结果图片保存路径(可选)
        show_result: 是否显示结果图像(可选)

        返回:
        result_image: 标记差异的结果图像
        difference_count: 差异区域数量
        """
        # 读取图像
        img1 = cv2.imread(img_path1)
        img2 = cv2.imread(img_path2)

        if img1 is None:
            raise FileNotFoundError(f"无法读取图像: {img_path1}")
        if img2 is None:
            raise FileNotFoundError(f"无法读取图像: {img_path2}")

        print(f"图像1尺寸: {img1.shape}, 图像2尺寸: {img2.shape}")

        # 找出差异
        differences, diff_map, M = self.find_differences(img1, img2)

        # 可视化结果
        result_img = self.visualize_differences(img1, img2, differences, M, output_path, show_result)

        # 保存结果
        self.last_result = {
            "result_image": result_img,
            "difference_map": diff_map,
            "difference_count": len(differences),
            "contours": differences,
            "image1": img1,
            "image2": img2,
            "homography_matrix": M,
            "image1_path": img_path1,
            "image2_path": img_path2
        }

        return result_img, len(differences)

    def save_difference_map(self, output_path):
        """保存差异图"""
        if self.last_result is None:
            raise ValueError("没有可用的结果，请先调用compare_images方法")

        cv2.imwrite(output_path, self.last_result["difference_map"])
        return output_path


def select_image():
    """通过图形界面选择图像路径"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    file_path = filedialog.askopenfilename(title="Select an image", filetypes=[("Image files", "*.png;*.jpg;*.jpeg")])
    return file_path


def main():
    # 通过图形界面选择图像路径
    print("Please select the first image")
    image1_path = select_image()
    print(f"Selected image 1: {image1_path}")

    print("Please select the second image")
    image2_path = select_image()
    print(f"Selected image 2: {image2_path}")

    # 创建比较器实例
    comparator = ImageComparator(min_contour_area=50, morph_kernel_size=3, ssim_threshold=0.85)

    try:
        # 比较图像
        result_img, diff_count = comparator.compare_images(image1_path, image2_path, output_path='result.png',
                                                           show_result=True)

        print(f"Found {diff_count} differences")
        print("Result saved to result.png")

        # 保存差异图
        comparator.save_difference_map("difference_map.png")
        print("Difference map saved to difference_map.png")

    except Exception as e:
        print(f"Error processing images: {str(e)}")


if __name__ == "__main__":
    main()




